# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo

import time

import coin.exchange.base.kr_rest.private_client_base as privcb
import coin.exchange.okex_futures.kr_rest.native_private_client as nprivc

from coin.exchange.base.kr_rest.rest_client_base import RestUpdate
from coin.exchange.okex_futures.kr_rest.currency import OkexCurrency
from coin.exchange.okex_futures.kr_rest.futures_constants import (
    active_contract_list,)
from coin.exchange.okex_futures.kr_rest.futures_product import (
    OkexFuturesProduct,)
from coin.exchange.okex_futures.kr_rest import util

import coin.proto.coin_order_enums_pb2 as coin_order
from coin.proto.coin_query_pb2 import (AccountBalance,
                                       CurrencyBalance,
                                       ProductPosition,
                                       MarginUsage,
                                       AccountPosition,
                                       ProductFillElement,
                                       ProductFillBundle,
                                       ProductOrderElement,
                                       ProductOrderBundle)
from coin.proto.coin_order_gateway_pb2 import (
    OrderSubmitResponseProto,
    OrderCancelResponseProto,
)


class OkexFuturesPrivateParser(object):
  @staticmethod
  def parse_balance(json_msg):
    """Balance json message.
    {
      "result": true,
      "info": {
        "btc": {
          "risk_rate": 10000,
          "account_rights": 0.09480204,
          "profit_unreal": 0,
          "profit_real": 0,
          "keep_deposit": 0
        },
        "btg": {
          "risk_rate": 10000,
          "account_rights": 0,
          "profit_unreal": 0,
          "profit_real": 0,
          "keep_deposit": 0
        },
    }
    """
    assert json_msg['result'] is True, json_msg
    currency_balance_list = []
    for currency, currency_info in json_msg['info'].items():
      # Sometimes OkexFutures adds new coins, just ignore these unknown coins.
      currency = OkexCurrency.FromStrNativeCurrencyNothrow(currency)
      if currency is None:
        continue

      total = float(currency_info['account_rights'])
      hold = float(currency_info['keep_deposit'])
      available = total - hold
      assert available >= 0.0, available
      assert hold >= 0.0, hold

      currency_balance = CurrencyBalance(currency=currency.currency,
                                         currency_native=currency.native_currency,
                                         available=available,
                                         hold=hold,
                                         total=total)
      currency_balance_list.append(currency_balance)

    return AccountBalance(
        exchange='Okex',
        market_type='Futures',
        each_balance=currency_balance_list,
    )

  @staticmethod
  def parse_margin(json_msg):
    assert json_msg['result'] is True, json_msg
    margin_usage_list = []
    for currency, currency_info in json_msg['info'].items():
      currency = OkexCurrency.FromStrNativeCurrencyNothrow(currency)
      if currency is None:
        continue

      total = float(currency_info['account_rights'])
      hold = float(currency_info['keep_deposit'])
      available = total - hold
      assert available >= 0.0, available
      assert hold >= 0.0, hold

      margin_usage = MarginUsage(currency=currency.currency,
                                 available_margin=available,
                                 hold_margin=hold,
                                 total_margin=total)
      margin_usage_list.append(margin_usage)
    return margin_usage_list

  @staticmethod
  def parse_position(json_msg):
    """
    {
      "result": true,
      "holding": [
        {
          "buy_price_avg": 10195.78064367,
          "symbol": "btc_usd",
          "lever_rate": 10,
          "buy_available": 46,
          "contract_id": 201803020000034,
          "buy_amount": 46,
          "buy_profit_real": 0.04824265,
          "contract_type": "this_week",
          "sell_amount": 308,
          "sell_price_cost": 10286.05555103,
          "buy_price_cost": 10195.78064367,
          "create_date": 1519564688000,
          "sell_price_avg": 10286.05555103,
          "sell_profit_real": -0.29010941,
          "sell_available": 308
        }
      ],
      "force_liqu_price": "20,283.42"
    }"""
    assert json_msg['result'] is True, json_msg
    liquidation_price = None
    if "force_liqu_price" in json_msg:
      liquidation_price_raw = json_msg['force_liqu_price']
      if type(liquidation_price_raw) == str:
        liquidation_price = float(liquidation_price_raw)
      elif type(liquidation_price_raw) == list:
        liquidation_price_total = sum([float(x.replace(',', '')) for x in liquidation_price_raw])
        liquidation_price = liquidation_price_total / len(liquidation_price_raw)
      else:
        pass

    product_position_list = []
    for position_info in json_msg['holding']:
      long_position = position_info['buy_amount']
      short_position = -position_info['sell_amount']
      net_position = long_position + short_position
      realized_pnl = (position_info['buy_profit_real'] + position_info['sell_profit_real'])
      product = OkexFuturesProduct.FromCurrencyPairAndContractType(position_info['symbol'],
                                                                   position_info['contract_type'])
      product_position = ProductPosition(symbol=product.symbol,
                                         native_symbol=product.native_symbol,
                                         long_position=long_position,
                                         short_position=short_position,
                                         net_position=net_position,
                                         realized_pnl=realized_pnl,
                                         liquidation_price=liquidation_price)
      product_position_list.append(product_position)
    return product_position_list

  # status field meaning:
  # -1 = cancelled,
  # 0 = unfilled,
  # 1 = partially filled,
  # 2 = fully filled,
  # 3 = ???  (status 3 is not mentioned in the doc, but observed)
  # 4 = cancel request in process
  # 5 = ??? (status 5 is not mentioned in the doc, but observed)
  @staticmethod
  def parse_fill(update_msg, product) -> ProductFillBundle:
    """Only first 50 filled orders are retrieved."""
    assert update_msg['result']
    currency_fill_list = []
    for update_elem in update_msg['orders']:
      assert update_elem['status'] in (-1, 0, 1, 2, 3, 4, 5), update_elem['status']
      if update_elem['status'] not in [1, 2]:
        continue
      if update_elem['contract_name'] != product.native_symbol:
        continue

      if update_elem['type'] == 1:
        side = coin_order.BUY_OPEN_ORDER
      elif update_elem['type'] == 2:
        side = coin_order.SELL_OPEN_ORDER
      elif update_elem['type'] == 3:
        side = coin_order.BUY_CLOSE_ORDER
      elif update_elem['type'] == 4:
        side = coin_order.SELL_CLOSE_ORDER
      else:
        raise ValueError("%s, %s" % (update_elem['type'], type(update_elem['type'])))

      currency_fill = ProductFillElement(symbol=product.symbol,
                                         side=side,
                                         price=float(update_elem['price']),
                                         qty=float(update_elem['deal_amount']),
                                         order_id=str(update_elem['order_id']))
      currency_fill_list.append(currency_fill)

    return ProductFillBundle(
        exchange='Okex',
        market_type='Futures',
        symbol=product.symbol,
        each_fill=currency_fill_list,
    )

  @staticmethod
  def parse_orders(update_msg, product):
    """Only first 50 filled orders are retrieved."""
    assert update_msg['result']
    currency_order_list = []
    for update_elem in update_msg['orders']:
      assert update_elem['status'] in (-1, 0, 1, 2, 3, 4, 5), update_elem['status']
      if update_elem['status'] not in [0, 1]:
        continue
      if update_elem['contract_name'] != product.native_symbol:
        continue

      if update_elem['type'] == 1:
        side = coin_order.BUY_OPEN_ORDER
      elif update_elem['type'] == 2:
        side = coin_order.SELL_OPEN_ORDER
      elif update_elem['type'] == 3:
        side = coin_order.BUY_CLOSE_ORDER
      elif update_elem['type'] == 4:
        side = coin_order.SELL_CLOSE_ORDER
      else:
        raise ValueError("%s, %s" % (update_elem['type'], type(update_elem['type'])))

      currency_order = ProductOrderElement(symbol=product.symbol,
                                           side=side,
                                           price=float(update_elem['price']),
                                           qty=float(update_elem['amount']),
                                           order_id=str(update_elem['order_id']),
                                           created_timestamp=int(
                                               float(update_elem['create_date']) * 1e+6))
      currency_order_list.append(currency_order)
    return ProductOrderBundle(exchange='Okex',
                              market_type='Futures',
                              symbol=product.symbol,
                              each_order=currency_order_list)

  @staticmethod
  def parse_submit_response(update_msg):
    if update_msg['result']:
      return OrderSubmitResponseProto(
          exchange='Okex',
          market_type='Futures',
          success=True,
          order_id=str(update_msg['order_id']),
      )
    else:
      return OrderSubmitResponseProto(exchange='Okex', market_type='Futures', success=False)

  @staticmethod
  def parse_cancel_response(update_msg):
    if update_msg['result']:
      return OrderCancelResponseProto(
          exchange='Okex',
          market_type='Futures',
          success=True,
          order_id=str(update_msg['order_id']),
      )
    else:
      return OrderCancelResponseProto(exchange='Okex', market_type='Futures', success=False)

  @staticmethod
  def summarize_cancel_all(update_list):
    success = all([update.msg.success for update in update_list])
    order_id_list = [update.msg.order_id for update in update_list]
    return OrderCancelResponseProto(
        exchange='Okex',
        market_type='Futures',
        success=success,
        order_id_list=order_id_list,
    )


class OkexFuturesPrivateClient(privcb.PrivateClientBase):
  ProductType = OkexFuturesProduct
  # 1: open long position
  # 2: open short position
  # 3: liquidate long position
  # 4: liquidate short position
  _trade_type_map = {
      'bid': 1,
      'ask': 2,
      'buy_open': 1,
      'sell_open': 2,
      'sell_close': 3,
      'buy_close': 4,
  }

  def __init__(self, key_file):
    privcb.PrivateClientBase.__init__(self)
    self.nprivc = nprivc.OkexFuturesNativePrivateClient(key_file=key_file)

  def query_account_balance_impl(self):
    # Product is not used. Always query all coin balances.
    update = self.nprivc.query_POST_futures_userinfo()
    update.msg = OkexFuturesPrivateParser.parse_balance(update.msg)
    return update

  def query_account_position_impl(self):
    margin_update = self.nprivc.query_POST_futures_userinfo()
    margin_update.msg = OkexFuturesPrivateParser.parse_margin(margin_update.msg)

    if True:
      # Query contract in active_contract_list one by one.
      holding_list = []
      force_liqu_price_list = []

      req_time = time.time()
      for currency_pair, contract_type in active_contract_list:
        update = self.nprivc.query_POST_futures_position(currency_pair, contract_type)
        assert update.msg['result']
        holding_list.extend(update.msg['holding'])
        force_liqu_price_list.append(update.msg['force_liqu_price'])
      res_time = time.time()

      raw_msg = {
          'result': True,
          'holding': holding_list,
          'force_liqu_price': force_liqu_price_list,
      }
      pos_update = RestUpdate(req_time, res_time, raw_msg)
      pos_update.msg = OkexFuturesPrivateParser.parse_position(pos_update.msg)
    else:
      assert isinstance(product, OkexFuturesProduct)
      currency_pair, contract_type = product.get_query_pair()
      pos_update = self.nprivc.query_POST_futures_position(currency_pair, contract_type)
      pos_update.msg = OkexFuturesPrivateParser.parse_position(pos_update.msg)

    account_position = AccountPosition(
        exchange='Okex',
        market_type='Futures',
        each_position=pos_update.msg,
        each_margin_usage=margin_update.msg,
    )
    native_msg = {
        "margin": margin_update.native_msg,
        "position": pos_update.native_msg,
    }
    update = RestUpdate(margin_update.req_timestamp,
                        pos_update.res_timestamp,
                        account_position,
                        native_msg=native_msg)
    return update

  def query_fills_impl(self, product):
    symbol, contract_type = product.get_query_pair()
    update = self.nprivc.query_POST_futures_order_info(symbol, contract_type, status='filled')
    update.msg = OkexFuturesPrivateParser.parse_fill(update.msg, product)
    return update

  def query_list_orders_impl(self, product):
    symbol, contract_type = product.get_query_pair()
    update = self.nprivc.query_POST_futures_order_info(symbol, contract_type, status='unfilled')
    update.msg = OkexFuturesPrivateParser.parse_orders(update.msg, product)
    return update

  def query_submit_impl(self, product, side, price, qty):
    # TODO add risk checker.
    assert price > 0, price
    assert qty > 0, qty
    price_str = util.fmt_price(price)

    update = self.nprivc.query_POST_futures_trade(product.order_symbol,
                                                  product.contract_type,
                                                  price_str,
                                                  qty,
                                                  self._trade_type_map[side])
    update.msg = OkexFuturesPrivateParser.parse_submit_response(update.msg)
    return update

  def query_cancel_impl(self, product, order_id):
    update = self.nprivc.query_POST_futures_cancel(product.order_symbol,
                                                   product.contract_type,
                                                   order_id)
    update.msg = OkexFuturesPrivateParser.parse_cancel_response(update.msg)
    return update

  # Okex does not provide cancel all API, emulate one.
  def query_cancel_product_impl(self, product):
    open_order_update = self.query_list_orders(product)
    open_orders = list(open_order_update.msg.each_order)
    if not open_orders:
      update = RestUpdate(time.time(), time.time(), None)
      update.msg = OrderCancelResponseProto(exchange='Okex', market_type='Futures', success=True)
      return update

    update_list = []
    for order in open_orders:
      single_update = self.query_cancel(product, order.order_id)
      update_list.append(single_update)
    update = RestUpdate(update_list[0].req_timestamp,
                        update_list[-1].res_timestamp, [upd.native_msg for upd in update_list])
    update.msg = OkexFuturesPrivateParser.summarize_cancel_all(update_list)
    return update
